"""
测试NSGA2算法 
"""
import matplotlib.pyplot as plt
from matplotlib.pyplot import axis
from NSGA2 import *
from function import *
from fitness import *


def main():
    nIter = 50
    nChr = 3
    nPop = 100
    pc = 0.6
    pm = 0.1
    etaC = 1
    etaM = 1
    func = function
    lb = -2
    rb = 2
    paretoPops, paretoFits = NSGA2(nIter, nChr, nPop, pc, pm, etaC, etaM, func, lb, rb)
    print(f"paretoFits适应度函数2维目标变量（100，2）: {paretoFits}")
    print(f"paretoPops种群3维决策变量（100，3）: {paretoPops}")

    # 理论最优解集合 
    x = np.linspace(-1 / np.sqrt(3), 1 / np.sqrt(3), 116).reshape(116, 1)
    X = np.concatenate((x, x, x), axis=1)
    thFits = fitness(X, function)

    plt.rcParams['font.sans-serif'] = 'KaiTi'  # 设置显示中文 
    fig = plt.figure(dpi=400)
    ax = fig.add_subplot(111)
    ax.plot(thFits[:, 0], thFits[:, 1], color='green', label='理论帕累托前沿')
    ax.scatter(paretoFits[:, 0], paretoFits[:, 1], color='red', label='实际解集')
    ax.legend()
    fig.savefig('test.png', dpi=400)

if __name__ == "__main__":
    main()
